# ------------------------------------------------------------------------------
# Calculate correlation per timepoint (and sample if stated) from the
# metacells (original or leiden) for different
# gene sets (split dependent on gene expression cutoff) for comparison with
# metacells (see corresponding files create_genesets.R,
# metacell_general_correlation_tp.R and eval_blueprint_genesets.R)
# Input: Seurat object, file with selected genes
# Output: files with correlation values (r-values and p-values)
# ------------------------------------------------------------------------------

library(Hmisc)
library(optparse)

#Parse arguments
option_list = list(
  make_option(c("-g","--selectedGenes"), 
              default="../../benchmark/celltypes/gene_expressed_over_hald_cells.txt",
              help="path to list with selected genes"),
  make_option(c("-m","--method"),
              default="metacell",
              help="method for metacell grouping (leiden[_SCT] or metacell)"),
  make_option(c("-s","--perSample"),action="store_true",
              default=FALSE,
              help="Shall the evaluation be done for each sample separatly"),
  make_option(c("-o","--outputFile"),
              default="timepoint_monocytes",
              help="Suffix of the output files")
)

opt_parser = OptionParser(option_list=option_list)
opt = parse_args(opt_parser)

pathSelectedGenes<-opt$selectedGenes
type<-opt$method
perSample<-opt$perSample
outputSuffix<-opt$outputFile

print(paste("Evaluating",type,"for gene set:"))
print(pathSelectedGenes)

print(paste("Evaluating each sample individually:", perSample))

#For leiden clustering
if(type=="leiden") {
  setwd("leiden_metacells/")
  pseudobulkFile<-"metacell_leiden.RDS"
  annotationFile<-"annotations_mc_leiden_tp.tsv"
  
  #Read pseudobulk data frame
  metacell.allsamples<-readRDS(pseudobulkFile)

  #For leiden clustering based on SCT counts
} else if(type=="leiden_SCT") {
  setwd("leiden_metacells/")
  pseudobulkFile<-"metacell_leiden_SCT.RDS"
  annotationFile<-"annotations_mc_leiden_SCT_tp.tsv"
  
  #Read pseudobulk data frame
  metacell.allsamples<-readRDS(pseudobulkFile)
  
} else if(type=="metacell"){
  setwd("metacell_general/metacell")
  
  metacellDir<-"metacells_K20_minCells10"
  setwd(metacellDir)

  pseudobulkFile<-"pseudobulk_metacell.RDS"
  annotationFile<-"annotations_metacell.tsv"
  
  metacell.allsamples<-readRDS(pseudobulkFile)
} else {
  stop("Metacell method type not known!")
}
  
##########################

#Read annotation data frame
annotations.allsamples<-read.table(annotationFile)
colnames(annotations.allsamples)[2]<-"timepoint"

#Select which genes shall be chosen for correlation (same as for single cell)
selected.genes<-read.table(pathSelectedGenes,
                           header=FALSE)
metacell.allsamples<-metacell.allsamples[selected.genes$V1,]

#Result data frame (correlation and pvalues)
corr.df<-NULL
pval.df<-NULL

correlationRes<-function(meta_counts,colName){
  
  #Be carefull: rcorr does not work with less than 5 samples
  corr.mc<-rcorr(t(meta_counts), type="spearman")
  
  #Create a pairwise data frame for the correlation
  corr.pairs.mc<-as.data.frame(as.table(corr.mc$r),
                               stringsAsFactors = FALSE)
  corr.pairs.mc<-corr.pairs.mc[corr.pairs.mc$Var1<corr.pairs.mc$Var2,]
  colnames(corr.pairs.mc)<-c("Gene1","Gene2",colName)
  
  #Create a pairwise data frame for the pvalue
  corr.pairs.pval<-as.data.frame(as.table(corr.mc$P),
                                 stringsAsFactors = FALSE)
  corr.pairs.pval<-corr.pairs.pval[corr.pairs.pval$Var1<corr.pairs.pval$Var2,]
  colnames(corr.pairs.pval)<-c("Gene1","Gene2",colName)
  
  return(list(corr.pairs.mc,corr.pairs.pval))
}

if(perSample){
  
  #Calculate correlation for each sample
  samples<-unique(annotations.allsamples$sample)
  for(sample in samples){
    print(paste("Processing sample:",sample))
    
    annot.sample<-annotations.allsamples[annotations.allsamples$sample==sample,]
    for(timepoint in unique(annot.sample$timepoint)){
      
      #Run the analysis only if at least 5 meta-cells exists
      #Probably increase the threshold again to more later ...
      mc.ids.timepoint<-annot.sample$metacell[annot.sample$timepoint==timepoint]
      if(length(mc.ids.timepoint)>4){
        
        print(paste("Calculate correlation for timepoint",timepoint))
        
        meta_counts<-metacell.allsamples[,mc.ids.timepoint]
        tmp<-correlationRes(meta_counts,colName = paste0(timepoint,"-",sample))
        corr.pairs.mc<-tmp[[1]]
        corr.pairs.pval<-tmp[[2]]
        
        #Concatinate the sample - timepoint pairs
        if(is.null(corr.df)){
          corr.df<-corr.pairs.mc
          pval.df<-corr.pairs.pval
        } else {
          corr.df<-merge(corr.df,corr.pairs.mc,by=c("Gene1","Gene2"),
                         all=TRUE)
          pval.df<-merge(pval.df,corr.pairs.pval,by=c("Gene1","Gene2"),
                         all=TRUE)
        }
        
      } else {
        print(paste("Skip timepoint",timepoint,"(too less metacells)"))
      }
    }
  }
  
} else {

  annot.sample<-annotations.allsamples
  for(timepoint in unique(annot.sample$timepoint)){
    
    #Run the analysis only if at least 5 meta-cells exists
    #Probably increase the threshold again to more later ...
    mc.ids.timepoint<-annot.sample$metacell[annot.sample$timepoint==timepoint]
    if(length(mc.ids.timepoint)>4){
      
      print(paste("Calculate correlation for timepoint",timepoint))
      
      meta_counts<-metacell.allsamples[,mc.ids.timepoint]
      tmp<-correlationRes(meta_counts,colName = timepoint)
      corr.pairs.mc<-tmp[[1]]
      corr.pairs.pval<-tmp[[2]]
      
      #Concatinate the sample - timepoint pairs
      if(is.null(corr.df)){
        corr.df<-corr.pairs.mc
        pval.df<-corr.pairs.pval
      } else {
        corr.df<-merge(corr.df,corr.pairs.mc,by=c("Gene1","Gene2"),
                       all=TRUE)
        pval.df<-merge(pval.df,corr.pairs.pval,by=c("Gene1","Gene2"),
                       all=TRUE)
      }
      
    } else {
      print(paste("Skip timepoint",timepoint,"(too less metacells)"))
    }
  }
}

write.table(corr.df,
            file=paste0("correlation_r_",outputSuffix,".tsv"),
            sep="\t",quote=FALSE)
write.table(pval.df,
            file=paste0("correlation_pval_",outputSuffix,".tsv"),
            sep="\t",quote=FALSE)